Le0v1n: autocast 是 torch.cuda.amp 模块中的一个类，它可以自动选择 GPU 操作的精度，以提高性能并保持精度。在使用 with autocast: 上下文管理器时，相关的操作（如 matmul，nn.Conv2d）会在计算过程中执行半精度运算。但是，最终计算的 loss 仍然是 float32 类型。 如果希望在使用 autocast 的同时进行梯度缩放，可以使用 torch.cuda.amp.GradScaler。它可以方便地执行梯度缩放的步骤，从而改善具有 float16 梯度的网络的收敛性。梯度缩放通过最小化梯度下溢来改善收敛性。 例如，在训练过程中，可以这样使用 autocast 和 GradScaler： [code=python] from torch.cuda.amp import autocast, GradScaler # 创建模型和优化器，默认精度为 float32 model = Net().cuda() optimizer = optim.SGD(model.parameters(), ...) # 在训练开始前创建一个 GradScaler 对象 scaler = GradScaler() for epoch in epochs: for input, target in data: optimizer.zero_grad() # 使用 autocast 进行前向传播 with autocast(): output = model(input) loss = loss_fn(output, target) # 缩放 loss 并进行反向传播 scaler.scale(loss).backward() # scaler.step() 首先对优化器分配的参数的梯度进行反缩放。如果这些梯度不包含 inf 或 NaN，则调optimizer.step()，否则跳过 optimizer.step() scaler.step(optimizer) # 更新缩放因子 scaler.update() [/code]